package cn.ycc1.mymiddle.chat;

import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.java.Log;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.Disposable;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.logging.StreamHandler;

/**
 * @author ycc
 * @date 2025/5/1
 */
@RestController
@RequestMapping("/chat-model")
public class ChatModelController {
    private final ChatModel chatModel;

    public ChatModelController(ChatModel chatModel) {
        this.chatModel = chatModel;
    }

    @RequestMapping("/chat")
    public String chat(String input) {
        ChatResponse response = chatModel.call(new Prompt(input));
        return response.getResult().getOutput().getText();
    }

    @RequestMapping("/chatWithPrompt")
    public String chatWithPrompt(String input) {
        Prompt prompt = new Prompt(input);
        ChatResponse response = chatModel.call(prompt);
        return response.getResult().getOutput().getText();
    }

    /**
     * 使用 SseEmitter：
     * 替代直接操作 HttpServletResponse，Spring 的 SseEmitter 专门为服务器发送事件设计
     * 自动处理异步响应生命周期
     * @param input
     * @return
     */
    @RequestMapping("/streamChat2")
    public SseEmitter streamChat2(String input) {
        SseEmitter emitter = new SseEmitter();

        Prompt prompt = new Prompt(input);
        Disposable subscription = chatModel.stream(prompt)
                .subscribe(
                        chatResponse -> {
                            try {
                                String content = new String(
                                        chatResponse.getResult().getOutput().getText().getBytes(StandardCharsets.UTF_8),
                                        StandardCharsets.UTF_8
                                );
                                emitter.send(SseEmitter.event()
                                        .data(content));
                            } catch (IOException e) {
                                throw new RuntimeException(e);
                            }
                        },
                        error -> emitter.completeWithError(error),
                        emitter::complete
                );

        emitter.onCompletion(subscription::dispose);
        emitter.onTimeout(() -> {
            subscription.dispose();
            emitter.complete();
        });

        return emitter;
    }

    // 重点：在@RequestMapping中指定produces
    @GetMapping(value = "/streamChat", produces = MediaType.TEXT_EVENT_STREAM_VALUE + ";charset=UTF-8")
    public SseEmitter streamChat(String input) {
        SseEmitter emitter = new SseEmitter(60_000L); // 设置超时时间60秒

        chatModel.stream(new Prompt(input))
                .subscribe(
                        chatResponse -> {
                            String content = chatResponse.getResult().getOutput().getText();
                            try {
                                emitter.send(SseEmitter.event()
                                        .data(content));
                            } catch (IOException e) {
                                throw new RuntimeException(e);
                            }
                        },
                        error -> emitter.completeWithError(error),
                        emitter::complete
                );

        // 处理客户端断开连接
        emitter.onCompletion(() -> System.out.println("SSE连接已完成"));
        emitter.onTimeout(() -> System.out.println("SSE连接超时"));

        return emitter;
    }
}
